Skip to content

P1-B: BufferHandle/BufferRef memory-model ABI + full user-facing Tensor→BufferRef cutover - #1453

Draft
YunjiQin wants to merge 7 commits into
hw-native-sys:mainfrom
YunjiQin:p1-b-buffer-handle-abi
Draft

P1-B: BufferHandle/BufferRef memory-model ABI + full user-facing Tensor→BufferRef cutover#1453
YunjiQin wants to merge 7 commits into
hw-native-sys:mainfrom
YunjiQin:p1-b-buffer-handle-abi

Conversation

@YunjiQin

@YunjiQin YunjiQin commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

P1-B: BufferHandle / BufferRef memory-model ABI + full user-facing Tensor→BufferRef cutover

Replaces the legacy "raw pointer + child_memory bool" L3+ memory mechanism with a typed,
self-describing BufferHandle / BufferRef ABI, flips the live L3→L2 (and L4→L3) dispatch wire from
Tensor blobs to BufferRef blobs, and makes the user-facing argument API uniformly BufferRef /
BufferHandle at every level — C++ Tensor survives only inside the L2 chip runtime.

User guide: docs/buffer-handle-abi.md. Frozen wire ABI:
.docs/L3-new/worker-memory-model/bufferhandle-abi.md. Phase plan: .docs/L3/P1.md.

Commits (read in order)

  1. Add the BufferHandle/BufferRef memory-model ABI + owner create_buffer — wire types
    (child_memoryaddress_space), versioned blob codec, owner BufferHandle + ImportRegistry
    (lazy map-once by canonical identity), self-describing refs, the view algebra
    (slice/transpose/permute/view/reshape), FORK_SHM COW base VA, Worker.create_buffer.
  2. Flip the live dispatch wire to BufferRef; three-way consumer split; device/remote/strided
    TaskArgs TensorBufferRef through the C++ dispatch; infer_deps keys on canonical identity;
    chip → materialize → run_from_blob, sub → MappedArg (torch.frombuffer), L4→L3 → re-export;
    DEVICE_MALLOC + REMOTE_SIDECAR + strided materialization.
  3. Resource APIs + ABI doc + cpput migration; remove TensorTaskArgs; L2 in-process materialization
    alloc_shared_tensor / alloc_child_tensor; migrate the C++ GoogleTest suite; fold TensorTaskArgs
    onto ChipStorageTaskArgs; uniform L2 in-process ref materialization.
  4. Move device-memory ownership from Orchestrator to Workermalloc/free/copy_to/copy_from/
    alloc_child_tensor live on the Worker (keyed by owner_worker_id); orch.* are thin forwards; the
    raw-pointer orch.malloc / orch.device_handle are removed.
  5. make_ref_arg + comm-domain VMM_WINDOW buffers; migrate examples & scene tests
    Worker.make_ref_arg / make_tensor_ref (memoized FORK_SHM ref over a pre-fork host tensor);
    allocate_domain hands back VMM_WINDOW BufferHandles (domain.buffers[name].ref); migrate every
    collective / demo / L2-L3 example and the scene-test framework.
  6. ChipStorageTaskArgs chip-only; delete create_host_buffer; sub scalars; orch.alloc→handle
    ChipStorageTaskArgs is the chip-only POD (ChipWorker._run_slot); Worker.run L2 takes a
    BufferRef TaskArgs/None; delete the create_host_buffer/MAP_HOST subsystem (superseded by
    create_buffer + lazy ref); sub callables receive scalars; handle.ref(dtype=) accepts a DataType
    enum; orch.alloc returns a BufferHandle over a single identity-keyed Orchestrator::alloc.

Verification

  • a2a3sim: allreduce collective ST corpus + L2/L3 examples (allreduce, domain_rank_map,
    dual_domain_overlap, ffn_tp_parallel, async/deferred notify, vector_add, create_buffer,
    l3_dependency/group, dynamic_register) — golden 0.
  • onboard a2a3 (Ascend910): L3→L2 chip dispatch, mixed host+device, L4→L3→sub re-export,
    sdma_async_completion.
  • pyut green: test_task_interface (109), test_callable_identity (79), test_sub_bufferref (8),
    test_l3_l2_message_queue + test_l3_l2_orch_comm (53), test_host_worker alloc/sub,
    test_scene_test_rehost, test_startup_readiness, test_remote_l3_lifecycle/protocol.
  • cpput 61/61 green (no-hardware subset: test_orchestrator, test_scheduler, test_tensormap,
    test_remote_wire/endpoint, …).
  • Every commit passes the pre-commit pyright / ruff / clang gate.

Remaining (follow-ups, not in this PR)

  • P2 (remote): the send side (RemoteTaskArgsWire) still carries Tensor metadata over the wire;
    full remote e2e; buffer lifecycle (release_buffer, in-flight retain / deferred-free). REMOTE_SIDECAR
    materialization stays reserved for P2.
  • Heap separation (SharedBufferArena): alloc_shared_tensor is the interim runtime-managed
    intermediate; the dedicated arena is a follow-up.

🤖 Generated with Claude Code

@coderabbitai

coderabbitai Bot commented Jul 24, 2026

Copy link
Copy Markdown

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: a3e9b77f-134b-42c7-8772-9b5ccdaa68da

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@YunjiQin
YunjiQin force-pushed the p1-b-buffer-handle-abi branch from ec3f15f to d82c15b Compare July 27, 2026 03:31
@YunjiQin YunjiQin changed the title P1-B: BufferHandle/BufferRef memory-model ABI + L3→L2 live cutover (phases B+C) P1-B: BufferHandle/BufferRef memory-model ABI + full user-facing Tensor→BufferRef cutover Jul 27, 2026
@YunjiQin
YunjiQin force-pushed the p1-b-buffer-handle-abi branch from d82c15b to 41cb93e Compare July 27, 2026 06:18
YunjiQin and others added 4 commits July 27, 2026 01:05
The typed, self-describing BufferHandle / BufferRef ABI that replaces the raw-pointer + child_memory
mechanism: wire types (child_memory→address_space), versioned blob codec (write/read_bufferref_blob),
owner-side BufferHandle + create_host_shared_buffer + ImportRegistry (lazy map-once by canonical
identity), self-describing refs (embedded descriptor, no eager handshake), the BufferRef view algebra
(slice/transpose/permute/view/reshape), FORK_SHM COW-inherited base VA, Worker.create_buffer, and
MAILBOX_SIZE 32768→40960 for the 272 B refs wire. Python buffer_handle.py mirrors buffer_handle.h
byte-for-byte (sizes pinned to the binding).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…evice/remote/strided

Flip TaskArgs Tensor→BufferRef end-to-end through the C++ dispatch; infer_deps keys a local ref by
canonical identity. Three-way consumer split (only the chip leaf rebuilds a C++ Tensor): chip →
materialize → run_from_blob; Python sub-worker → MappedArg (torch.frombuffer); nested L4→L3 → re-export
(no map on the forwarding hop, canonical identity invariant across the edge — frozen model §5/§8: an
L4-owned backing forwarded L4→L3→L2 keeps ONE identity at all three layers, so deps key on identity not
VA and aliases/retain-release do not split). Device-memory path (DEVICE_MALLOC materialize/wrap),
REMOTE_SIDECAR backend-rewrite at add_ref, and strided-ref → strided-Tensor materialization. Proven
L3→L2 + L4→L3→sub in sim and onboard a2a3.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… L2 in-process materialization

C-phase resource APIs (Worker.alloc_shared_tensor, Orchestrator.alloc_child_tensor); docs/
buffer-handle-abi.md (identity invariant across every edge; re-export preserves it); migrate the C++
GoogleTest orchestrator/scheduler/remote/child_memory tests to the BufferRef ABI; remove TensorTaskArgs
(fold onto ChipStorageTaskArgs); restore device-pointer provenance over the BufferRef wire; relax
create_buffer to L2 and materialize BufferRef args in-process on an L2 run (uniform user API at every
level).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Worker owns and manages the device-memory lifecycle (malloc / free / copy_to / copy_from /
alloc_child_tensor, keyed by owner_worker_id); the Orchestrator's memory APIs become thin forwards. The
raw-pointer orch.malloc / orch.device_handle are removed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@YunjiQin
YunjiQin force-pushed the p1-b-buffer-handle-abi branch 3 times, most recently from 3e6cde6 to cc65ce6 Compare July 27, 2026 12:40
YunjiQin and others added 3 commits July 27, 2026 05:40
… scene tests to BufferRef/handle

Worker.make_ref_arg + simpler_setup.make_tensor_ref name a pre-fork host tensor as a memoized FORK_SHM
ref; alloc_shared_tensor becomes the alloc→BufferRef managed intermediate; comm-domain windows are
exposed as VMM_WINDOW BufferHandles (domain.buffers[name].ref). Migrate the L2 examples, task_timing st,
child_memory, and the scene-test L3 rehost path / L3 group + dynamic_register scene tests / post-fork
host-buffer ST onto create_buffer + BufferRef args; delete the dead _rewrite_blob_host_addrs; sync
docs/comm-domain.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…alars; orch.alloc→handle

Complete the Tensor removal from the user-facing path:
- ChipStorageTaskArgs is chip-only (ChipWorker._run_slot); Worker.run at L2 takes a BufferRef TaskArgs
  or None (legacy passthrough removed); migrate the remaining wire-TaskArgs pyut + remote
  materialization + test_callable_identity + scene-rehost UT fakes off add_tensor.
- Delete the create_host_buffer / MAP_HOST subsystem (superseded by create_buffer + lazy ref).
- handle.ref(dtype=) accepts a DataType enum; Python sub callables receive scalars (MappedArgs).
- orch.alloc returns a BufferHandle over a single identity-keyed Orchestrator::alloc; finish Tensor
  removal from the L3-L2 message queue / region.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two O(1) warm-path static checks per the frozen model (worker-memory-model §1.4/§6, bufferhandle-abi
§3/§4.1), no ABI/wire change:

- Capability matrix: BufferHandleDescriptor.__post_init__ rejects an unsupported
  address_space×backend_kind (the four illegal combos HOST×VMM_WINDOW / HOST×DEVICE_MALLOC /
  DEVICE×FORK_SHM / DEVICE×POSIX_SHM). Runs on both owner construction and wire decode, so a bad combo
  fails before dispatch and cannot ride the wire. REMOTE_SIDECAR rows stay allowed here (its P1 blanket
  reject is ImportRegistry.materialize's job).
- access⊆granted: TaskArgs.add_ref (C++ binding) rejects an arg whose TensorArgType needs access the
  handle does not grant (INPUT→READ, OUTPUT_EXISTING→WRITE, INOUT→READWRITE). Catches e.g. a READ-only
  FORK_SHM COW buffer tagged OUTPUT_EXISTING (a forked child's writes would silently not reach the
  parent).

Not implemented here (already covered / P2, per the task doc): device→owner-chip topology
(_child_prov_check_dispatch), overlapping-alias (infer_deps serializes by identity), unknown
abi/backend/REMOTE_SIDECAR reject (already), stale-generation/ABA/visibility (P2).

Tests: test_descriptor_rejects_bad_capability_combo / accepts_legal_combos;
TestAddRefAccessSubset. Verified a2a3sim async_notify golden 0 (guards don't misfire on legit flows).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant